import logging
from pathlib import Path
from traceback import format_exc
from typing import Union

import torch
import torchaudio
from demucs.pretrained import get_model
from demucs.audio import AudioFile
from demucs.apply import apply_model

# ✅ 支持的格式集中管理
SUPPORTED_FORMATS = {"wav", "flac", "ogg", "mp3"}


def save_audio(path: str, waveform: torch.Tensor, sample_rate: int, fmt: str):
    """
    Save audio in different formats using appropriate backends:
    - soundfile for wav/flac/ogg
    - pydub (ffmpeg) for mp3/aac/m4a
    """
    import soundfile as sf
    from pydub import AudioSegment
    import numpy as np
    logging.info(f"Saving audio to {path} in {fmt} format")
    fmt = fmt.lower()
    path = str(path)

    # Ensure waveform is 2D
    if waveform.dim() == 1:
        waveform = waveform.unsqueeze(0)  # [C, T]

    # Convert to numpy for external libraries
    waveform_np = waveform.cpu().numpy()

    if fmt in {"wav", "flac", "ogg"}:
        # Use soundfile
        waveform_np = waveform_np.T  # soundfile expects shape [T, C]
        subtype = "PCM_16" if fmt in {"wav", "flac"} else None  # ogg = float32
        sf.write(path, waveform_np, sample_rate,
                 format=fmt.upper(), subtype=subtype)

    elif fmt in {"mp3", "aac", "m4a"}:
        # Use pydub (requires ffmpeg)
        # Convert to 16-bit mono/stereo PCM for export
        waveform_int16 = (waveform_np * 32767.0).clip(-32768,
                                                      32767).astype(np.int16)
        audio = AudioSegment(
            waveform_int16.T.tobytes(),
            frame_rate=sample_rate,
            sample_width=2,
            channels=waveform_np.shape[0],
        )
        audio.export(path, format=fmt)

    else:
        raise ValueError(f"Unsupported output format: {fmt}")

    logging.debug(f"Saved: {path} (format={fmt})")


def run_demucs(path: Union[str, Path], output_format="wav") -> Path:
    """
    Runs Demucs source separation on the given file path.

    :param path: Input audio file path
    :param output_format: Output audio format (e.g., 'wav', 'mp3', etc.)
    :return: Path to the output directory
    """
    try:
        path = Path(path)
        fmt = output_format.lower()

        # ✅ 格式校验
        if fmt not in SUPPORTED_FORMATS:
            raise ValueError(
                f"Unsupported format '{fmt}'. Supported formats: {', '.join(SUPPORTED_FORMATS)}")

        model = get_model(name="htdemucs")
        model.eval()
        device = "cuda" if torch.cuda.is_available() else "cpu"
        model.to(device)

        af = AudioFile(path)
        wav = af.read(
            streams=slice(0, 1),
            samplerate=model.samplerate,
            channels=model.audio_channels
        ).to(device)

        if wav.dim() == 1:
            wav = wav.unsqueeze(0)
        elif wav.dim() == 2:
            wav = wav.unsqueeze(0)  # Shape: (1, C, T)

        sources = apply_model(model, wav, device=device)[0]

        out_dir = path.with_name(path.stem + f"_{fmt}_demucs")
        out_dir.mkdir(parents=True, exist_ok=True)

        # ✅ 统一保存逻辑，无论是否 dict
        if isinstance(sources, dict):
            items = sources.items()
        else:
            items = zip(model.sources, sources)

        for name, stem in items:
            stem = stem.detach().cpu()
            save_audio(str(out_dir / f"{name}.{fmt}"),
                       stem, model.samplerate, fmt)

        logging.info(f"Separation done. Output saved to {out_dir}")
        return out_dir

    except Exception as e:
        logging.error(f"Failed to separate audio: {e}\n{format_exc()}")
        raise
